#!/usr/bin/env python3

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Command line utility for reading and writing Avro files."""

import csv
import functools
import json
import optparse
import os.path
import sys

import avro
import avro.datafile
import avro.errors
import avro.io
import avro.schema

_AVRO_DIR = os.path.abspath(os.path.dirname(avro.__file__))


def _version():
    with open(os.path.join(_AVRO_DIR, 'VERSION.txt')) as v:
        return v.read()


_AVRO_VERSION = _version()


def print_json(row):
    print(json.dumps(row))


def print_json_pretty(row):
    """Pretty print JSON"""
    # Need to work around https://bugs.python.org/issue16333
    # where json.dumps leaves trailing spaces.
    result = json.dumps(row, sort_keys=True, indent=4).replace(' \n', '\n')
    print(result)


_write_row = csv.writer(sys.stdout).writerow


def print_csv(row):
    # We sort the keys to the fields will be in the same place
    # FIXME: Do we want to do it in schema order?
    _write_row([row[key] for key in sorted(row)])


def select_printer(format):
    return {
        "json": print_json,
        "json-pretty": print_json_pretty,
        "csv": print_csv
    }[format]


def record_match(expr, record):
    return eval(expr, None, {"r": record})


def parse_fields(fields):
    fields = fields or ''
    if not fields.strip():
        return None

    return [field.strip() for field in fields.split(',') if field.strip()]


def field_selector(fields):
    fields = set(fields)

    def keys_filter(obj):
        return dict((k, obj[k]) for k in (set(obj) & fields))
    return keys_filter


def print_avro(avro, opts):
    if opts.header and (opts.format != "csv"):
        raise avro.errors.UsageError("--header applies only to CSV format")

    # Apply filter first
    if opts.filter:
        predicate = functools.partial(record_match, opts.filter)
        avro = (r for r in avro if predicate(r))

    for i in range(opts.skip):
        try:
            next(avro)
        except StopIteration:
            return

    fields = parse_fields(opts.fields)
    if fields:
        fs = field_selector(fields)
        avro = (fs(r) for r in avro)

    printer = select_printer(opts.format)
    for i, record in enumerate(avro):
        if i == 0 and opts.header:
            _write_row(sorted(record.keys()))
        if i >= opts.count:
            break
        printer(record)


def print_schema(avro):
    schema = avro.schema
    # Pretty print
    print(json.dumps(json.loads(schema), indent=4))


def cat(opts, args):
    if not args:
        raise avro.errors.UsageError("No files to show")
    for filename in args:
        with avro.datafile.DataFileReader(open(filename, 'rb'), avro.io.DatumReader()) as avro_:
            if opts.print_schema:
                print_schema(avro_)
                continue
            print_avro(avro_, opts)


def _open(filename, mode):
    if filename == "-":
        return {
            "rb": sys.stdin,
            "wb": sys.stdout
        }[mode]

    return open(filename, mode)


def iter_json(info, _):
    for i in info:
        try:
            s = i.decode()
        except AttributeError:
            s = i
        yield json.loads(s)


def convert(value, field):
    type = field.type.type
    if type == "union":
        return convert_union(value, field)

    return  {
        "int": int,
        "long": int,
        "float": float,
        "double": float,
        "string": str,
        "bytes": bytes,
        "boolean": bool,
        "null": lambda _: None,
        "union": lambda v: convert_union(v, field),
    }[type](value)


def convert_union(value, field):
    for name in [s.name for s in field.type.schemas]:
        try:
            return convert(name)(value)
        except ValueError:
            continue


def iter_csv(info, schema):
    header = [field.name for field in schema.fields]
    for row in csv.reader((getattr(i, "decode", lambda: i)() for i in info)):
        values = [convert(v, f) for v, f in zip(row, schema.fields)]
        yield dict(zip(header, values))


def guess_input_type(files):
    if not files:
        return None

    ext = os.path.splitext(files[0])[1].lower()
    if ext in (".json", ".js"):
        return "json"
    elif ext in (".csv",):
        return "csv"

    return None


def write(opts, files):
    if not opts.schema:
        raise avro.errors.UsageError("No schema specified")

    input_type = opts.input_type or guess_input_type(files)
    if not input_type:
        raise avro.errors.UsageError("Can't guess input file type (not .json or .csv)")
    iter_records = {"json": iter_json, "csv": iter_csv}[input_type]

    try:
        with open(opts.schema) as schema_file:
            schema = avro.schema.parse(schema_file.read())
        out = _open(opts.output, "wb")
    except (IOError, OSError) as e:
        raise avro.errors.UsageError("Can't open file - %s" % e)

    writer = avro.datafile.DataFileWriter(getattr(out, 'buffer', out), avro.io.DatumWriter(), schema)

    for filename in (files or ["-"]):
        info = _open(filename, "rb")
        for record in iter_records(info, schema):
            writer.append(record)

    writer.close()


def main(argv):
    parser = optparse.OptionParser(description="Display/write for Avro files",
                                   version=_AVRO_VERSION,
                                   usage="usage: %prog cat|write [options] FILE [FILE...]")
    # cat options
    cat_options = optparse.OptionGroup(parser, "cat options")
    cat_options.add_option("-n", "--count", default=float("Infinity"),
                           help="number of records to print", type=int)
    cat_options.add_option("-s", "--skip", help="number of records to skip",
                           type=int, default=0)
    cat_options.add_option("-f", "--format", help="record format",
                           default="json",
                           choices=["json", "csv", "json-pretty"])
    cat_options.add_option("--header", help="print CSV header", default=False,
                           action="store_true")
    cat_options.add_option("--filter", help="filter records (e.g. r['age']>1)",
                           default=None)
    cat_options.add_option("--print-schema", help="print schema",
                           action="store_true", default=False)
    cat_options.add_option('--fields', default=None,
                           help='fields to show, comma separated (show all by default)')
    parser.add_option_group(cat_options)

    # write options
    write_options = optparse.OptionGroup(parser, "write options")
    write_options.add_option("--schema", help="schema file (required)")
    write_options.add_option("--input-type",
                             help="input file(s) type (json or csv)",
                             choices=["json", "csv"], default=None)
    write_options.add_option("-o", "--output", help="output file", default="-")
    parser.add_option_group(write_options)

    opts, args = parser.parse_args(argv[1:])
    if len(args) < 1:
        parser.error("You much specify `cat` or `write`")  # Will exit

    command_name = args.pop(0)
    try:
        command = {
            "cat": cat,
            "write": write,
        }[command_name]
    except KeyError:
        raise avro.errors.UsageError("Unknown command - {!s}".format(command_name))
    command(opts, args)


if __name__ == "__main__":
    main(sys.argv)
